"""
可视化工具模块

提供结果可视化功能
"""

import numpy as np
import matplotlib.pyplot as plt
from typing import Dict, List, Tuple, Optional
from pathlib import Path

class ResultVisualizer:
    """结果可视化类"""
    
    def __init__(self, results: Dict):
        """
        初始化可视化器
        
        参数:
            results (Dict): 计算结果字典
        """
        self.results = results
        
    def plot_axial_profiles(self, save_dir: Optional[str] = None):
        """
        绘制轴向分布图
        
        参数:
            save_dir (Optional[str]): 图片保存目录
        """
        # 创建保存目录
        if save_dir:
            Path(save_dir).mkdir(parents=True, exist_ok=True)
            
        # 绘制压力分布
        self._plot_pressure_profile(save_dir)
        
        # 绘制温度分布
        self._plot_temperature_profile(save_dir)
        
        # 绘制空泡率分布
        self._plot_void_fraction_profile(save_dir)
        
        # 绘制速度分布
        self._plot_velocity_profile(save_dir)
        
    def _plot_pressure_profile(self, save_dir: Optional[str]):
        """绘制压力分布"""
        plt.figure(figsize=(10, 6))
        plt.plot(self.results['axial_position'], self.results['pressure'])
        plt.xlabel('轴向位置 (m)')
        plt.ylabel('压力 (Pa)')
        plt.title('轴向压力分布')
        plt.grid(True)
        
        if save_dir:
            plt.savefig(Path(save_dir) / 'pressure_profile.png')
        plt.close()
        
    def _plot_temperature_profile(self, save_dir: Optional[str]):
        """绘制温度分布"""
        plt.figure(figsize=(10, 6))
        plt.plot(self.results['axial_position'], self.results['temperature'])
        plt.xlabel('轴向位置 (m)')
        plt.ylabel('温度 (K)')
        plt.title('轴向温度分布')
        plt.grid(True)
        
        if save_dir:
            plt.savefig(Path(save_dir) / 'temperature_profile.png')
        plt.close()
        
    def _plot_void_fraction_profile(self, save_dir: Optional[str]):
        """绘制空泡率分布"""
        plt.figure(figsize=(10, 6))
        plt.plot(self.results['axial_position'], self.results['void_fraction'])
        plt.xlabel('轴向位置 (m)')
        plt.ylabel('空泡率')
        plt.title('轴向空泡率分布')
        plt.grid(True)
        
        if save_dir:
            plt.savefig(Path(save_dir) / 'void_fraction_profile.png')
        plt.close()
        
    def _plot_velocity_profile(self, save_dir: Optional[str]):
        """绘制速度分布"""
        plt.figure(figsize=(10, 6))
        plt.plot(self.results['axial_position'], self.results['velocity_axial'], 
                label='轴向速度')
        plt.plot(self.results['axial_position'], self.results['velocity_lateral'], 
                label='横向速度')
        plt.xlabel('轴向位置 (m)')
        plt.ylabel('速度 (m/s)')
        plt.title('轴向速度分布')
        plt.legend()
        plt.grid(True)
        
        if save_dir:
            plt.savefig(Path(save_dir) / 'velocity_profile.png')
        plt.close() 